import numpy as np
import cv2
import csv
import pandas
import argparse
import pcl, pcl.pcl_visualization

#timestamp, temperature, horizantal angle (in degree), vertical angle (in degree), distance, intensity, x, y, z
scan_data_type = np.dtype(dict(names=['timestamp', 'temperature', 'hangle', 'vangle', 'dist', 'intensity' 'x','y','z'],
    formats=['f8','f4','f4','f4', 'f4', 'f4', 'f4', 'f4', 'f4'],
    offsets=[0, 8, 12, 16, 20, 24, 28, 32, 36], itemsize=40))
    
def read_scan(path):    
    with open(path, 'r') as f:
        f_csv = csv.reader(f)
        scan_data = [[float(item) for item in row] for row in f_csv]
        return np.float64(scan_data)

def transformSPH2XYZ(hangles, vangles, _row):
    row = np.zeros((_row.shape[0], 3))
    a = np.deg2rad(hangles)  # 水平转角
    b = np.deg2rad(vangles)  # 竖直转角
    r = row  # 距离
    row[:, 0] = r * np.sin(b) * np.cos(a)
    row[:, 1] = -r * np.sin(b) * np.sin(a)
    row[:, 2] = r * np.cos(b)
    return row

def denoise_float32(array):
    kernel = np.array((
        [0.0625, 0.125, 0.0625],
        [0.125, 0.25, 0.125],
        [0.0625, 0.125, 0.0625]), dtype=np.float32)
    
    denoised = cv2.filter2D(array, -1, kernel)
    return denoised

def polar_filter(data, angular_resolution, filter2=None):
    width = int(360 / angular_resolution)
    height = int(180 / angular_resolution)
    range_image = np.zeros((width, height, 6), dtype=np.float32)
    count = np.zeros((width, height), dtype=np.int32)
    h, v = np.mod(np.int32(data[:, 2] / angular_resolution), width), np.mod(np.int32(data[:, 3] / angular_resolution), height)
    for i in range(data.shape[0]):
        range_image[h[i], v[i], :] = range_image[h[i], v[i], :] + data[i, [4, 1, 5, 6, 7, 8]]
        count[h[i], v[i]] = count[h[i], v[i]] + 1
    
    count_idx = count.nonzero()
    for i in range(6):
        range_image[count_idx[0], count_idx[1], i] = range_image[count_idx[0], count_idx[1], i] / count[count_idx[0], count_idx[1]]
    
    if filter2 is not None:
        range_image = filter2(range_image)

    # horizantal angle (in degree), vertical angle (in degree), distance, intensity, x, y, z
    # filtered_data = np.zeros((len(polar), 3), dtype=np.float32)
    # polar = np.array(list(polar))
    # filtered_data[:, [0, 1]] = polar * angular_resolution
    # filtered_data[:, 2] = range_image[polar[:, 0], polar[:, 1], 0]
    # filtered_data[:, 2] = range_image[polar[:, 0], polar[:, 1], 0]
    # filtered_data[:, 3] = range_image[polar[:, 0], polar[:, 1], 1]
    # filtered_data[:, 4] = range_image[polar[:, 0], polar[:, 1], 2]
    # # xyzi[:, 5] = range_image[polar[:, 0], polar[:, 1], 0]
    # filtered_data[:, 5:] = transformSPH2XYZ(filtered_data[:, :3])
    # return range_image, filtered_data
    # return transformSPH2XYZ(filtered_data)
    return range_image

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', default='./data/fixed/fixed.csv')
    # parser.add_argument('-o', '--output', default='./data/fixed/fixed_polar.pcd')
    parser.add_argument('-r', '--resolution', default=0.2)
    parser.add_argument('-d', '--denoise', default=False)
    args = parser.parse_args()
    data = read_scan(args.input)
    filter = None
    if args.denoise: filter = denoise_float32
    data1 = data[data[:, 3] <= 0, :]
    data2 = data[data[:, 3] >= 0, :]
    dri1 = polar_filter(data1, args.resolution, filter2=filter)
    intensity_img = dri1[:, :, 2]
    intensity_img = np.uint8(intensity_img * 255 / intensity_img.max())
    cv2.imwrite(args.input + '.1-%s-%s.jpg' % (dri1.shape[0], dri1.shape[1]), intensity_img)
    data1 = dri1[:, :, 3:].reshape((-1, 3))
    pc = pcl.PointCloud(data1)
    output = args.input + '.1.pcd'
    pcl.save(pc, output, format='pcd', binary=True)

    dri2 = polar_filter(data2, args.resolution, filter2=filter)
    intensity_img = dri2[:, :, 2]
    intensity_img = np.uint8(intensity_img * 255 / intensity_img.max())
    cv2.imwrite(args.input + '.2-%s-%s.jpg' % (dri2.shape[0], dri2.shape[1]), intensity_img)
    data2 = dri2[:, :, 3:].reshape((-1, 3))
    pc = pcl.PointCloud(data2)
    output = args.input + '.2.pcd'
    pcl.save(pc, output, format='pcd', binary=True)
    # viewer = pcl.pcl_visualization.CloudViewing()
    # viewer.ShowMonochromeCloud(pc)
    # new_data = pandas.DataFrame(dxyz, columns=['hangle', 'vangle', 'dist', 'temperature', 'intensity', 'x', 'y', 'z'])
    # new_data.to_csv(args.output)    